FastSAMで高速道路を走行する車両をセグメンテーションをしてみました（上下線や追い越し車線も判定してみました）
1 はじめに
CX 事業本部 delivery部の平内（SIN）です。
Meta社による Segment Anything Model（SAM)は、セグメンテーションのための汎用モデルで、ファインチューニングなしで、あらゆる物体がセグメンテーションできます。
そして、それを爆速にしたのがFastSAM(Fast Segment Anything)です。
FastSAMは、SAM がトレーニングされた SA-1B データセットの 2% でトレーニングされたモデルということで、SAMよりも 50倍高速に実行されるとアナウンスされています。 手元でも試してみましたが、SAMだと数秒かかっていた処理が、100msぐらいで終わってしまうのを確認できました。
モデルのサイズも、SAMが、2.5 GByteだったの対して、FastSAMでは、145 MByteと小さくなっており、取り回しも軽くなっていると思います。
$ ls -la -rw-r--r-- 1 root root 144943063 Jun 20 16:01 FastSAM.pt -rw-r--r-- 1 root root 2564550879 Apr 4 15:56 sam_vit_h_4b8939.pth
当然と言えば当然ですが、SAMに比べると、やや検出は甘いですが、使い所を考えれば、充分に機能できると思います。
ということで、今回は、高速道路を撮影した動画で、リアルタイムなセグメンテーションを試して見ました。
最初に、試してみた様子をご覧ください。
元動画は24FPSでしたが、１フレーム毎に読み飛ばしたので、12FPSで動作している感じです。 セグメンテーションされた車両は、「上り」「下り」「追い越し車線」ごとに色を変えて輪郭を縁取っています。また、側道を走る車は、対象外としました。
2 環境
(1) nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3
作業した環境は、Jetson AGX Orin 開発者キットです。
Jetsonでは、NVIDIAで配布されているPyTorch Container for Jetson and JetPackを使用すると、CUDAでPyTorchを利用する環境が簡単に準備できます。
$ sudo docker pull nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3 $ sudo docker run -it --rm --runtime nvidia --network host nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3 # python3 -c "import torch;print(torch.cuda.is_available())" True
上記をベースに、必要なライブラリをインストールしたイメージを作業用に作成しました。なお、FastSAMのrequirements.txtでは、OpenCVの最新版が指定されているのですが、Dockerイメージにあらかじめセットアップされているものと競合してしまうので、一旦、無効化しています。
DockerFile
FROM nvcr.io/nvidia/l4t-pytorch:r35.2.1-pth2.0-py3 RUN apt-get update RUN apt-get install -y python3-tk # pipでcv2がインストールされた場合に競合してしまうので無効化する RUN mv /usr/lib/python3.8/dist-packages/cv2 /usr/lib/python3.8/dist-packages/cv2.bak RUN mv /usr/local/lib/python3.8/dist-packages/cv2 /usr/local/lib/python3.8/dist-packages/cv2. RUN pip install matplotlib pyyaml tqdm pandas psutil scipy seaborn gitpython RUN pip install ultralytics
$ sudo docker build -t fast_sam:latest .
作成した作業用イメージは、以下の感じで使用しています。
docker-start.sh
xhost + sudo docker run -it --rm --runtime nvidia --shm-size=1g -v $(pwd)/home:/home -e DISPLAY=:0 --network host fast_sam:latest
(2) Setup
FastSAMのセットアップです。
# cd home
// FastSAM
# git clone https://github.com/CASIA-IVA-Lab/FastSAM.git
# pip -q install -r FastSAM/requirements.txt
// dependencies
# pip -q install roboflow supervision jupyter_bbox_widget
// model
# mkdir -p ./weights
# wget -P ./weights -q https://huggingface.co/spaces/An-619/FastSAM/resolve/main/weights/FastSAM.pt
# tree . -L 2
.
├── index.py <= 今回作成したプログラムです。
├── judgement_util.py <= 衝突検知用の関数です。
├── FastSAM
│ ├── app_gradio.py
│ ├── assets
│ ├── cog.yaml
│ ├── examples
│ ├── fastsam
│ ├── images
│ ├── Inference.py
│ ├── LICENSE
│ ├── MORE_USAGES.md
│ ├── output
│ ├── predict.py
│ ├── README.md
│ ├── requirements.txt
│ ├── segpredict.py
│ ├── setup.py
│ └── utils
└── weights
└─── FastSAM.pt
3 車線の判定
高速道路の上下線や、追い越し車線を走行している車両の判別は、以下のような手順で行なっています。
まず、FastSAMでは、confやiouのパラメータ指定に応じて、画面上のオブジェクトが検出されています。
そこで、道路上に「上り線」「下り線」及び、それぞれの「追い越し車線」の直線を引きます。
次に、検出されたマスクと車線の衝突検知を行っています。
検出されたマスクは、輪郭点の集合となっていますが、そのままでは、ちょっと計算量が大きくなってしまうので、一旦、外接する矩形に変換し、矩形と直線の衝突検知としています。
衝突検知は、下記で公開されているものを利用させてい頂きました。
凸多角形と直線の交差判定/交点
judgement_util.py
# 凸多角形と直線の交差判定/交点
# https://tjkendev.github.io/procon-library/python/geometry/line_convex_polygon_intersection.html
def find_zero(x0, x1, f):
v0 = f(x0)
if v0 == 0:
return x0 + 1
left = x0
right = x1 + 1
while left + 1 < right:
mid = (left + right) >> 1
if v0 * f(mid) >= 0:
left = mid
else:
right = mid
return right
def binary_search(f, L):
left = 0
right = L
while left + 1 < right:
mid = (left + right) >> 1
if f(mid) < 0:
left = mid
else:
right = mid
return right
def line_polygon_intersection(p0, p1, qs):
x0, y0 = p0
x1, y1 = p1
dx = x1 - x0
dy = y1 - y0
h = lambda p: (p[0] - x0) * dy - (p[1] - y0) * dx
L = len(qs)
i0 = i1 = -1
v0 = h(qs[0])
v1 = h(qs[L - 1])
if v0 == v1:
v2 = h(qs[1])
# assert v0 != v2
if v0 < v2:
i0 = L - 1
else:
i1 = L - 1
else:
v2 = h(qs[1])
if v1 > v0 <= v2:
i0 = 0
elif v1 < v0 >= v2:
i1 = 0
else:
g = lambda x: min((v1 - v0) * x / (L - 1) + v0, h(qs[x]))
i0 = binary_search(lambda x: g(x + 1) - g(x), L - 1)
if i1 == -1:
B = i0 - L
k = binary_search(lambda x: h(qs[B + x]) - h(qs[B + x + 1]), L)
i1 = (i0 + k) % L
else:
B = i1 - L
k = binary_search(lambda x: h(qs[B + x + 1]) - h(qs[B + x]), L)
i0 = (i1 + k) % L
if h(qs[i0]) * h(qs[i1]) > 0:
# a line and a polygon are disjoint
return []
# a vertex to the left side of a line: i0
# a vertex to the right side of a line: i1
f = lambda i: h(qs[i - L])
k0 = find_zero(i1, i0 if i1 < i0 else i0 + L, f) % L
k1 = find_zero(i0, i1 if i0 < i1 else i1 + L, f) % L
# vertices to the left side of a line: k0, k0+1, ..., k1-2, k1-1
# vertices to the right side of a line: k1, k1+1, ..., k0-2, k0-1
if k0 == k1:
return [k0]
return [k0, k1]
# e_i is a line segment between v_{i-1} and v_i
# p0, p1 = (0, 0), (1, 1)
# qs = [(0, 2), (2, 0), (4, 0), (4, 2), (2, 4), (0, 4)]
# print(line_polygon_intersection(p0, p1, qs))
# # => "[4, 1]": cross points is on either e_4 or e_1
# p0, p1 = (0, -1), (4, 0)
# qs = [(0, 2), (2, 0), (4, 0), (4, 2), (2, 4), (0, 4)]
# # => "[3]": cross point is on e_3
最終的に出力される映像は、以下のようになります。車線として定義されていない側道の車両は、結果的に無視されることになります。
4 コード
実装したコードです。
import os
import cv2
import torch
import numpy as np
import FastSAM.fastsam as fastsam
import judgement_util
def collision_detection(line, rect):
result = judgement_util.line_polygon_intersection(line[0], line[1], rect)
return len(result) > 0
def main():
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("DEVICE:{}".format(DEVICE))
FAST_SAM_CHECKPOINT = "./weights/FastSAM.pt"
print("FAST_SAM_CHECKPOINT:{}".format(FAST_SAM_CHECKPOINT))
model = fastsam.FastSAM(FAST_SAM_CHECKPOINT)
cap = cv2.VideoCapture("./highway.mp4")
if cap.isOpened() == False:
print("cv2.VideoCapture() faild.")
fps = cap.get(cv2.CAP_PROP_FPS)
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
print("FPS:{} WIDTH:{} HEIGHT:{}".format(fps, width, height))
# 道路のライン
y0 = 50 # 道路の上限
y1 = 360 # 道路の下限
up_0 = [(280, y0), (120, y1)] # 上り左車線
up_1 = [(310, y0), (240, y1)] # 上り右車線（追い越し）
down_0 = [(405, y0), (605, y1)] # 下り左車線
down_1 = [(375, y0), (485, y1)] # 下り右車線（追い越し）
counter = 0
while cap.isOpened():
ret, frame = cap.read()
counter += 1
if ret == True:
if counter % 2 == 0:
everything_results = model(
frame,
device=DEVICE,
retina_masks=True,
imgsz=1024,
conf=0.8,
iou=0.9,
)
annotations = everything_results[0].masks.data
try:
print("All annotations :{}".format(len(annotations)))
except Exception as e:
print("EROR: {}".format(e))
continue
annotations = annotations.cpu().numpy()
inference_frame = frame.copy()
contour_up_0 = []
contour_up_1 = []
contour_down_0 = []
contour_down_1 = []
for i, mask in enumerate(annotations):
annotation = mask.astype(np.uint8)
contours, hierarchy = cv2.findContours(
annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
)
for contour in contours:
# 外接の矩形を求める
x, y, w, h = cv2.boundingRect(contour)
if h >= 60: # 車のサイズとして異常値を排除する（車線等が検出されてしまうため）
continue
rect = [(x, y), (x, y + h), (x + w, y + h), (x + w, y)]
# 道路ラインとの衝突検知
if collision_detection(up_0, rect):
contour_up_0.append(contour)
if collision_detection(up_1, rect):
contour_up_1.append(contour)
if collision_detection(down_0, rect):
contour_down_0.append(contour)
if collision_detection(down_1, rect):
contour_down_1.append(contour)
cv2.drawContours(inference_frame, contour_up_0, -1, (255, 150, 0), 2)
cv2.drawContours(inference_frame, contour_up_1, -1, (255, 0, 0), 2)
cv2.drawContours(inference_frame, contour_down_0, -1, (0, 150, 255), 2)
cv2.drawContours(inference_frame, contour_down_1, -1, (0, 0, 255), 2)
cv2.imshow("Org", frame)
cv2.imshow("Inference", inference_frame)
if cv2.waitKey(25) & 0xFF == ord("q"):
break
else:
break
cap.release()
cv2.destroyAllWindows()
main()
5 最後に
今回は、Segment Anythingの高速版であるFastSAMで、動画をセグメンテーションしてみました。
以前にも書きましたが、Segment Anythingは、Computer Visionにおけるゲームチェンジャーだと思っていますが、処理速度の制限で、応用範囲が限られると感じていました。FastSAMの登場で、一気に新しい応用範囲が広がったと感じています。
